#ifndef LERP_N_D_H
#define LERP_N_D_H

#include "kernel_operator.h"

namespace Lerp {
using namespace AscendC;

constexpr int32_t FP32_BUFFER_NUM = 2;
constexpr int32_t FP16_BUFFER_NUM = 1;
constexpr int32_t BYTE_BLOCK = 32;
constexpr int32_t FP16_TMP_BUF_COUNT = 5 * 2;
constexpr int32_t FP32_TMP_BUF_COUNT = 1;

constexpr float POS_ONE = 1.0;
constexpr float NEG_ONE = -1.0;
constexpr float POS_HALF = 0.5;

template <typename T>
class LerpND {
public:
    __aicore__ inline LerpND() = default;
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR end, GM_ADDR weight, GM_ADDR y,
                                const LerpTilingData* __restrict tilingData);
    __aicore__ inline void Process();

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyIn(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void Compute(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void CopyOut(int64_t gmOffset, int64_t dataCount);

private:
    TPipe pipe;
#if ORIG_DTYPE_START == DT_FLOAT
    TQue<QuePosition::VECIN, FP32_BUFFER_NUM> startQue, endQue, weightQue;
    TQue<QuePosition::VECOUT, FP32_BUFFER_NUM> yQue;
#else
    TQue<QuePosition::VECIN, FP16_BUFFER_NUM> startQue, endQue, weightQue;
    TQue<QuePosition::VECOUT, FP16_BUFFER_NUM> yQue;
#endif
    TBuf<QuePosition::VECCALC> tempValBuf;
    GlobalTensor<T> startGM, endGM, weightGM, yGM;
    LocalTensor<float> tempValLT;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;

    // tiling params
    uint64_t ubMaxProcCount = 0;
    uint64_t totalDataCount = 0;
    uint64_t loopTime = 0;
    uint64_t tailCount = 0;
};

template <typename T>
__aicore__ inline void LerpND<T>::Init(GM_ADDR start, GM_ADDR end, GM_ADDR weight, GM_ADDR y,
                                       const LerpTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);
    ubMaxProcCount = tilingData->ubMaxProcCount;
    totalDataCount = tilingData->totalDataCount;
    loopTime = tilingData->loopTime;
    tailCount = tilingData->tailCount;

    startGM.SetGlobalBuffer((__gm__ T*)start, totalDataCount);
    endGM.SetGlobalBuffer((__gm__ T*)end, totalDataCount);
    weightGM.SetGlobalBuffer((__gm__ T*)weight, totalDataCount);
    yGM.SetGlobalBuffer((__gm__ T*)y, totalDataCount);

    uint64_t singleBufferSize = ubMaxProcCount * sizeof(T);
#if ORIG_DTYPE_START == DT_FLOAT
    pipe.InitBuffer(startQue, FP32_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(endQue, FP32_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(weightQue, FP32_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(yQue, FP32_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(tempValBuf, singleBufferSize * FP32_TMP_BUF_COUNT);
#else
    pipe.InitBuffer(startQue, FP16_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(endQue, FP16_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(weightQue, FP16_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(yQue, FP16_BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(tempValBuf, singleBufferSize * FP16_TMP_BUF_COUNT);
#endif
    tempValLT = tempValBuf.Get<float>();
}

template <typename T>
__aicore__ inline void LerpND<T>::Process() {
    int64_t gmOffset = 0;
    for (int64_t i = 0; i < loopTime; i++) {
        CopyIn(gmOffset, ubMaxProcCount);
        Compute(gmOffset, ubMaxProcCount);
        CopyOut(gmOffset, ubMaxProcCount);
        gmOffset += ubMaxProcCount;
    }
    if (tailCount) {
        int64_t alignDataCount = CeilAlignA2B(tailCount, perBlockCount);
        CopyIn(gmOffset, alignDataCount);
        Compute(gmOffset, alignDataCount);
        CopyOut(gmOffset, alignDataCount);
    }
}

template <typename T>
__aicore__ inline void LerpND<T>::CopyIn(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> startInLT = startQue.AllocTensor<T>();
    LocalTensor<T> endInLT = endQue.AllocTensor<T>();
    LocalTensor<T> weightInLT = weightQue.AllocTensor<T>();
    DataCopy(startInLT, startGM[gmOffset], dataCount);
    DataCopy(endInLT, endGM[gmOffset], dataCount);
    DataCopy(weightInLT, weightGM[gmOffset], dataCount);
    startQue.EnQue(startInLT);
    endQue.EnQue(endInLT);
    weightQue.EnQue(weightInLT);
}

template <typename T>
__aicore__ inline void LerpND<T>::Compute(int64_t gmOffset, int64_t dataCount) {
#if ORIG_DTYPE_START == DT_FLOAT
    LocalTensor<float> startLT = startQue.DeQue<float>();
    LocalTensor<float> endLT = endQue.DeQue<float>();
    LocalTensor<float> weightLT = weightQue.DeQue<float>();
    LocalTensor<float> yLT = yQue.AllocTensor<float>();
#else
    LocalTensor<half> startFP16 = startQue.DeQue<half>();
    LocalTensor<half> endFP16 = endQue.DeQue<half>();
    LocalTensor<half> weightFP16 = weightQue.DeQue<half>();
    LocalTensor<half> yFP16 = yQue.AllocTensor<half>();
    LocalTensor<float> startLT = tempValLT[ubMaxProcCount];
    LocalTensor<float> endLT = tempValLT[ubMaxProcCount * 2];
    LocalTensor<float> weightLT = tempValLT[ubMaxProcCount * 3];
    LocalTensor<float> yLT = tempValLT[ubMaxProcCount * 4];
    Cast(startLT, startFP16, RoundMode::CAST_NONE, dataCount);
    Cast(endLT, endFP16, RoundMode::CAST_NONE, dataCount);
    Cast(weightLT, weightFP16, RoundMode::CAST_NONE, dataCount);
#endif
    LocalTensor<uint8_t> mask = tempValLT.ReinterpretCast<uint8_t>();

    Duplicate(yLT, POS_HALF, dataCount);
    Compare(mask, weightLT, yLT, CMPMODE::LT, dataCount);
    Adds(yLT, weightLT, NEG_ONE, dataCount);
    Select(yLT, mask, yLT, weightLT, SELMODE::VSEL_CMPMASK_SPR, dataCount);
    Select(weightLT, mask, endLT, startLT, SELMODE::VSEL_CMPMASK_SPR, dataCount);
    Muls(startLT, startLT, NEG_ONE, dataCount);
    Add(endLT, endLT, startLT, dataCount);
    Mul(endLT, endLT, yLT, dataCount);
    Add(yLT, endLT, weightLT, dataCount);

#if ORIG_DTYPE_START == DT_FLOAT
    yQue.EnQue(yLT);
    startQue.FreeTensor(startLT);
    endQue.FreeTensor(endLT);
    weightQue.FreeTensor(weightLT);
#else
    Cast(yFP16, yLT, RoundMode::CAST_RINT, dataCount);
    yQue.EnQue(yFP16);
    startQue.FreeTensor(startFP16);
    endQue.FreeTensor(endFP16);
    weightQue.FreeTensor(weightFP16);
#endif
}

template <typename T>
__aicore__ inline void LerpND<T>::CopyOut(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> yOutLT = yQue.DeQue<T>();
    DataCopy(yGM[gmOffset], yOutLT, dataCount);
    yQue.FreeTensor(yOutLT);
}
}  // namespace Lerp

#endif  // LERP_N_D_H